/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include "parcsr_iluamg_preconditioner.h"
#include "parcsr_strength.h"
#include "parcsr_coarsen.h"
#include "parcsr_interpolation.h"
#include "parcsr_rap.h"
#include "parcsr_sparsification.h"
#include "parcsr_smooth.h"
#include "parcsr_restriction.h"
#include "parcsr_prolongation.h"
#include "parcsr_matvec.h"
#include "parcsr_transpose.h"
#include "parcsr_blockjacobiilu_preconditioner.h"

namespace YHAMG
{

class ParCSR_ILUAMG_Preconditioner::Smoother : public ParCSR_BlockJacobiILU_Preconditioner
{
public:
	void operator()(const ParCSR_Matrix& A, const Par_Vector& b, const Par_Vector& x, const Par_Vector& w) const
	{
		ParCSR_MatVec(-1.0, A, x, 1.0, b, w);
		Apply(w, w);
		x.AddScaled(1.0, w);
	}
};

void ParCSR_ILUAMG_Preconditioner::Strength(const ParCSR_Matrix& A, bool* LocLoc_s, bool* LocExt_s) const
{
	double threshold = StrengthThreshold;
	ParCSR_Strength(A, LocLoc_s, LocExt_s, threshold);
}

void ParCSR_ILUAMG_Preconditioner::Coarsening(const ParCSR_Matrix& A, const bool* LocLoc_s, const bool* LocExt_s, int* CF_map, int AGGRESSIVE) const
{
	if (AGGRESSIVE)
	{
		if (CoarsenType == 0)
			ParCSR_Aggressive_HMIS_Coarsening(A, LocLoc_s, LocExt_s, CF_map);
		else if (CoarsenType == 1)
			ParCSR_Aggressive_PMIS_Coarsening(A, LocLoc_s, LocExt_s, CF_map);
	}
	else
	{
		if (CoarsenType == 0)
			ParCSR_HMIS_Coarsening(A, LocLoc_s, LocExt_s, CF_map);
		else if (CoarsenType == 1)
			ParCSR_PMIS_Coarsening(A, LocLoc_s, LocExt_s, CF_map);
	}
}

void ParCSR_ILUAMG_Preconditioner::Interpolation(const ParCSR_Matrix& A, const bool* LocLoc_s, const bool* LocExt_s, int* CF_map, ParCSR_Matrix& P, int AGGRESSIVE) const
{
	int iterations = AGGRESSIVE ? 3 : 1;
	int min_elements = InterpMinElements;
	int max_elements = InterpMaxElements;
	double truncation_factor = TruncationFactor;
	ParCSR_Interpolation(A, LocLoc_s, LocExt_s, CF_map, P, iterations, min_elements, max_elements, truncation_factor);
}

void ParCSR_ILUAMG_Preconditioner::RAP(const ParCSR_Matrix& R, const ParCSR_Matrix& A, const ParCSR_Matrix& P, ParCSR_Matrix& C) const
{
	ParCSR_RAP(R, A, P, C);
}

void ParCSR_ILUAMG_Preconditioner::Sparsification(const ParCSR_Matrix& A, ParCSR_Matrix& B) const
{
	double threshold = SparsificationThreshold;
	ParCSR_Sparsification(A, B, threshold);
}

void ParCSR_ILUAMG_Preconditioner::Smooth(int level, const Par_Vector& b, const Par_Vector& x) const
{
	S[level](A[level], b, x, r[level]);
}

void ParCSR_ILUAMG_Preconditioner::Restriction(int level, const Par_Vector& b, const Par_Vector& x, const Par_Vector& g) const
{
	ParCSR_Restriction(R[level], A[level], b, x, g, r[level]);
}

void ParCSR_ILUAMG_Preconditioner::Prolongation(int level, const Par_Vector& e, const Par_Vector& x) const
{
	ParCSR_Prolongation(P[level], e, x);
}

void ParCSR_ILUAMG_Preconditioner::V_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const
{
	if (level == num_levels - 1)
	{
		for (int i = 0; i < CoarseSweeps; ++i)
			Smooth(level, b, x);
	}
	else
	{
		for (int i = 0; i < PreSweeps; ++i)
			Smooth(level, b, x);

		Restriction(level, b, x, g[level]);

		e[level].Fill(0.0);

		V_Cycle(level + 1, g[level], e[level]);

		Prolongation(level, e[level], x);

		for (int i = 0; i < PostSweeps; ++i)
			Smooth(level, b, x);
	}
}

void ParCSR_ILUAMG_Preconditioner::W_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const
{
	if (level == num_levels - 1)
	{
		for (int i = 0; i < CoarseSweeps; ++i)
			Smooth(level, b, x);
	}
	else
	{
		for (int i = 0; i < PreSweeps; ++i)
			Smooth(level, b, x);

		Restriction(level, b, x, g[level]);

		e[level].Fill(0.0);

		W_Cycle(level + 1, g[level], e[level]);
		W_Cycle(level + 1, g[level], e[level]);

		Prolongation(level, e[level], x);

		for (int i = 0; i < PostSweeps; ++i)
			Smooth(level, b, x);
	}
}

void ParCSR_ILUAMG_Preconditioner::F_Cycle(int level, const Par_Vector& b, const Par_Vector& x) const
{
	if (level == num_levels - 1)
	{
		x.Fill(0.0);
		for (int i = 0; i < CoarseSweeps; ++i)
			Smooth(level, b, x);
	}
	else
	{
		ParCSR_MatVec(1.0, R[level], b, 0.0, g[level], g[level]);

		F_Cycle(level + 1, g[level], e[level]);

		ParCSR_MatVec(1.0, P[level], e[level], 0.0, x, x);

		V_Cycle(level, b, x);
	}
}

ParCSR_ILUAMG_Preconditioner::ParCSR_ILUAMG_Preconditioner()
	: num_levels(0),
	A(0),
	P(0),
	R(0),
	S(0),
	r(0),
	e(0),
	g(0),
	MaxLevels(20),
	CoarseSize(8),
	StrengthThreshold(0.25),
	AggNumLevels(0),
	CoarsenType(0),
	InterpMinElements(4),
	InterpMaxElements(4),
	TruncationFactor(0.1),
	SparsificationThreshold(0.0),
	CycleType(0),
	PreSweeps(1),
	PostSweeps(1),
	CoarseSweeps(2),
	PrintStats(0)
{
}

ParCSR_ILUAMG_Preconditioner::ParCSR_ILUAMG_Preconditioner(
	int    max_levels,
	int    coarse_size,
	double strength_threshold,
	int    agg_num_levels,
	int    coarsen_type, 
	int    interp_min_elements,
	int    interp_max_elements,
	double truncation_factor,
	double sparsification_threshold,
	int    cycle_type,
	int    pre_sweeps,
	int    post_sweeps,
	int    coarse_sweeps,
	int    print_stats)
	: num_levels(0),
	A(0),
	P(0),
	R(0),
	S(0),
	r(0),
	e(0),
	g(0),
	MaxLevels(max_levels),
	CoarseSize(coarse_size),
	StrengthThreshold(strength_threshold),
	AggNumLevels(agg_num_levels),
	CoarsenType(coarsen_type),
	InterpMinElements(interp_min_elements),
	InterpMaxElements(interp_max_elements),
	TruncationFactor(truncation_factor),
	SparsificationThreshold(sparsification_threshold),
	CycleType(cycle_type),
	PreSweeps(pre_sweeps),
	PostSweeps(post_sweeps),
	CoarseSweeps(coarse_sweeps),
	PrintStats(print_stats)
{
}

ParCSR_ILUAMG_Preconditioner::~ParCSR_ILUAMG_Preconditioner()
{
	if (A) delete[] A;
	if (P) delete[] P;
	if (R) delete[] R;
	if (S) delete[] S;
	if (r) delete[] r;
	if (e) delete[] e;
	if (g) delete[] g;
}

void ParCSR_ILUAMG_Preconditioner::Free()
{
	if (A) delete[] A;
	if (P) delete[] P;
	if (R) delete[] R;
	if (S) delete[] S;
	if (r) delete[] r;
	if (e) delete[] e;
	if (g) delete[] g;
	num_levels = 0;
	A = 0;
	P = 0;
	R = 0;
	S = 0;
	r = 0;
	e = 0;
	g = 0;
}

void ParCSR_ILUAMG_Preconditioner::Setup(const ParCSR_Matrix& _A, int REUSE)
{
	if (REUSE == 0)
	{
		Free();

		comm = _A.comm;

		int rank;
		MPI_Comm_rank(comm, &rank);

		A = new ParCSR_Matrix[MaxLevels];
		P = new ParCSR_Matrix[MaxLevels - 1];
		R = new ParCSR_Matrix[MaxLevels - 1];
		r = new Par_Vector[MaxLevels];
		e = new Par_Vector[MaxLevels - 1];
		g = new Par_Vector[MaxLevels - 1];
		S = new Smoother[MaxLevels];

		long global_size0[2];
		double grid_complexity = 0.0;
		double operator_complexity = 0.0;

		if (PrintStats && rank == 0)
		{
			std::cout << "Multigrid Hierarchy:\n";
			std::cout << "Level\tRows\tNz\n";
			std::cout << "----------------------------------------------------------------\n";
		}

		A[0].Reference(_A);
		A[0].SetupHalo();

		int level = 0;

		while (1)
		{
			S[level].Setup(A[level], 0);

			int n = A[level].LocLoc.size[0];

			long local_size[2];
			long global_size[2];

			local_size[0] = n;
			local_size[1] = A[level].LocLoc.rowptr[n] + A[level].LocExt.rowptr[n];

			MPI_Allreduce(local_size, global_size, 2, MPI_LONG, MPI_SUM, comm);

			if (level == 0)
			{
				global_size0[0] = global_size[0];
				global_size0[1] = global_size[1];
			}

			grid_complexity += (double)global_size[0] / global_size0[0];
			operator_complexity += (double)global_size[1] / global_size0[1];

			if (PrintStats && rank == 0)
				std::cout << level << "\t" << global_size[0] << "\t" << global_size[1] << "\n";

			r[level].comm = comm;
			r[level].Resize(n);

			if (level == MaxLevels - 1 || global_size[0] <= (long)CoarseSize) break;

			bool* LocLoc_s = new bool[A[level].LocLoc.rowptr[n]];
			bool* LocExt_s = new bool[A[level].LocExt.rowptr[n]];
			int* CF_map = new int[n];

			Strength(A[level], LocLoc_s, LocExt_s);

			Coarsening(A[level], LocLoc_s, LocExt_s, CF_map, level < AggNumLevels);
			Interpolation(A[level], LocLoc_s, LocExt_s, CF_map, P[level], level < AggNumLevels);

			delete[] LocLoc_s;
			delete[] LocExt_s;
			delete[] CF_map;

			P[level].SetupHalo();
			
			ParCSR_Transpose(P[level], R[level]);

			R[level].SetupHalo();

			ParCSR_Matrix C;
			RAP(R[level], A[level], P[level], C);
			Sparsification(C, A[level +1]);

			A[level + 1].SetupHalo();

			e[level].comm = comm;
			g[level].comm = comm;
			e[level].Resize(A[level + 1].LocLoc.size[0]);
			g[level].Resize(A[level + 1].LocLoc.size[0]);

			++level;
		}

		num_levels = level + 1;
		
		if (PrintStats && rank == 0)
		{
			std::cout << "----------------------------------------------------------------\n";
			std::cout << "Number of Levels: " << num_levels << "\n";
			std::cout << "Grid Complexity: " << grid_complexity << "\n";
			std::cout << "Operator Complexity: " << operator_complexity << "\n";
		}
	}
	else
	{
		A[0].Reference(_A);
		A[0].SetupHalo();
		S[0].Setup(A[0], 1);
	}
}

int ParCSR_ILUAMG_Preconditioner::InSize() const
{
	return A[0].LocLoc.size[0];
}

int ParCSR_ILUAMG_Preconditioner::OutSize() const
{
	return A[0].LocLoc.size[0];
}

void ParCSR_ILUAMG_Preconditioner::Apply(const Par_Vector& b, const Par_Vector& x) const
{
	switch (CycleType)
	{
	case 0:
		x.Fill(0.0);
		V_Cycle(0, b, x);
		break;
	case 1:
		x.Fill(0.0);
		W_Cycle(0, b, x);
		break;
	case 2:
		F_Cycle(0, b, x);
		break;
	}
}

}